package org.ala.tiktools.sharddb.tx;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadPoolExecutor;

import org.ala.tiktools.logs.LoggerFactory;
import org.ala.tiktools.mybatis.spring.TikToolsMyBatisSpringContext;
import org.ala.tiktools.sharddb.db.ThreadLocalLookupKey;
import org.ala.tiktools.sharddb.tx.ShardDataSourceTransactionManager.ShardDataSourceTransactionObject;
import org.apache.ibatis.session.SqlSessionFactory;
import org.slf4j.Logger;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionException;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.TransactionSynchronizationManager;

/**
 *	分布式事务协调者（本地分布式事务）
 *	<p>	兼容原非分布式机制
 *	<p>	每个ShardTransactionStatus代表子事务参与者
 *
 *	<p>	预处理代码包装成task任务
 *	<p>		持有runnable预处理代码体
 *	<p>	协调者commit时，所有的预处理任务被扔进线程池执行
 *	<p>		此时创建ShardTransactionStatus，connection连接绑定线程环境
 *	<p>		预处理执行完，线程解绑，但不归还connection
 *	<p>		执行MyBatis绑定的前置动作（真实执行提交的线程已经不是跑预提交的线程了，所以要提前执行）
 *	<p>		解了绑的线程归还给线程池
 *	<p>	所有参与者预处理代码都跑完时，执行整体提交
 *	<p>		此时执行每个ShardTransactionStatus的commit
 *	<p>		执行完成后归还connection
 *	<p>	预处理阶段如果有task抛出异常
 *	<p>		执行所有task的rollback
 *	<p>		执行完成后归还connection
 *
 *  @author ala
 *  @date 2024-12-27 14:53
 */
public class ShardTransaction implements IShardTransaction {


    static Logger log = LoggerFactory.getLogger("tiktools", "shardTx");


    /**
     *  事务超时时间
     */
    protected long timeout;
    /**
     *  事务当前状态
     *  <p> 进行中 - 已提交 - 回滚
     */
    protected ShardTransactionStatusEnums status = ShardTransactionStatusEnums.Running;
    /**
     *  事务管理器
     */
    protected ShardDataSourceTransactionManager transactionManager;
    /**
     * 事务定义
     */
    protected TransactionDefinition transactionDefinition;
    /**
     * 事务任务队列
     */
    protected LinkedList<TransactionTask> tasks = new LinkedList<>();
    /**
     * 执行任务线程池
     */
    protected ThreadPoolExecutor threadPool;


    @Override
    public void withTxAutoCommit(String name, String shardKey, Runnable run)  throws Exception {
        withTxAutoCommit(name, shardKey, () -> {run.run(); return null;});
    }

    @Override
    public <T> T withTxAutoCommit(String name, String shardKey, Callable<T> call)  throws Exception {
        //  只有运行中的事务可以执行
        if (!ShardTransactionStatusEnums.Running.equals(status)) {
            throw new ShardTxException("只有运行中的事务允许执行");
        }

        //  设置分库环境
        ThreadLocalLookupKey.setLookup(shardKey);
        TransactionStatus ts = transactionManager.getTransaction(transactionDefinition);
        try {
            T res = call.call();
            transactionManager.commit(ts);
            return res;
        } catch (Exception e) {
            log.error(String.format("[ShardTransaction withTxAutoCommit] 事务操作出现异常 name:%s", name), e);
            transactionManager.rollback(ts);
            throw e;
        } finally {
            ThreadLocalLookupKey.clear();
        }
    }

    @Override
    public  void withTx2PCSync(String name, String shardKey, Runnable run) throws Exception {
        withTx2PCSync(name, shardKey, () -> {run.run(); return null;});
    }

    @SuppressWarnings("unchecked")
	@Override
    public <T> T withTx2PCSync(String name, String shardKey, Callable<T> run) throws Exception {
        //  只有运行中的事务可以执行
        if (!ShardTransactionStatusEnums.Running.equals(status)) {
            throw new ShardTxException("只有运行中的事务允许执行");
        }
        //  往执行队列里new一个任务
        TransactionTask task = TransactionTask.build(name, shardKey, null, run);
        doTransactionTask(task);
        //  如果执行过程中产生异常，往外抛。异常信息包装在
        if (task.e != null) {
            log.error(String.format("[ShardTransaction withTxAutoCommit] 事务操作出现异常 name:%s", name), task.e);
            rollback(null);
            throw task.e;
        }
        return (T) task.res;
    }
    @Override
    public void withTx2PCAsync(String name, String shardKey, Object metadata, Runnable run) throws Exception {
        //  只有运行中的事务可以执行
        if (!ShardTransactionStatusEnums.Running.equals(status)) {
            throw new ShardTxException("只有运行中的事务允许执行");
        }
        //  只是new一个任务，放入队列中，commit之前再一并提交
        tasks.add(TransactionTask.build(name, shardKey, null, () -> {run.run(); return null;}).metadata(metadata));
    }
    

    @Override
    public void commit(ICompensation compensation) throws Exception {
        //  只有运行中的事务允许提交
        if (!ShardTransactionStatusEnums.Running.equals(status)) {
            throw new ShardTxException("只有运行中的事务允许提交");
        }

        status = ShardTransactionStatusEnums.Commited;
        //  检测任务队列中是否还有没提交的任务，若有，扔线程池执行。若没有，提交每一个任务
        List<TransactionTask> _tasks = new LinkedList<>();
        tasks.forEach(t -> {
            if (TransactionTask.Status.UnBegin.equals(t.status)) {
                _tasks.add(t);
            }
        });

        //  等待全部任务执行完毕，并且没有错误
        CountDownLatch latch = new CountDownLatch(_tasks.size());
        boolean allSubmit = true;
        for (TransactionTask t : _tasks) {
            try {
                threadPool.execute(() -> {
                	try {
                		doTransactionTask(t);
                	} finally {
                		latch.countDown();
                	}
                });
            } catch (Exception e) {
                //  提交线程池失败，做回滚处理
                t.e = e;
                latch.countDown();
                allSubmit = false;
                break;
            }
        }
        //  检测是否存在任务执行超时
        if (allSubmit) {
        	latch.await();
//            boolean succ = latch.await(timeout, TimeUnit.MILLISECONDS);
//            if (!succ) {
//                //  给还没执行完的任务超时异常
//                for (TransactionTask t : tasks) {
//                    if (TransactionTask.Status.Running.equals(t.status)) {
//                        t.e = new TimeoutException("任务执行超时");
//                        log.error(String.format("[ShardTransaction withTxAutoCommit] 事务操作出现超时 name:%s", t.name), t.e);
//                        rollback(compensation);
//                        throw t.e;
//                    }
//                }
//            }
        }

        //  检测是否全部预提交成功
        for (TransactionTask t : tasks) {
            if (!TransactionTask.Status.Success.equals(t.status)) {
                log.error(String.format("[ShardTransaction withTxAutoCommit] 事务操作出现异常 name:%s", t.name), t.e);
                rollback(compensation);
                throw t.e;
            }
        }
        //  提交全部任务
        for (TransactionTask t : tasks) {
            try {
                transactionManager.commit(t.ts);
            } catch (TransactionException e) {
                log.error(String.format("[ShardTransaction withTxAutoCommit] 事务commit出现异常 name:%s", t.name), t.e);
                rollback(compensation);
                throw e;
            } finally {
            	
            }
        }
    }
    
    
    /**
     * 	执行预提交任务
     */
    protected void doTransactionTask(TransactionTask task) {
    	try {
    		ThreadLocalLookupKey.setLookup(task.shardKey);
    		task.ts = (ShardTransactionStatus) transactionManager.getTransaction(transactionDefinition);
    		task.doTask();
    		
    		//	初始化事务环境
    		task.ts.setContext(new ShardTransactionContext());
    	} catch (Exception e) {
    		e.printStackTrace();
    	} finally {
    		//	注册清理Mybatis资源任务
    		((ShardDataSourceTransactionObject)task.ts.getTransaction()).appendSynchronizations(new ShardSqlSessionSynchronization(task.ts.getContext()));
    		//	清空事务环境，线程回收给其他人用
    		TransactionSynchronizationManager.clear();
    		TransactionSynchronizationManager.unbindResource(((ShardTransactionStatus) task.ts).getDataSource());
    		TransactionSynchronizationManager.unbindResource(TikToolsMyBatisSpringContext.getBean(SqlSessionFactory.class));
    		ThreadLocalLookupKey.clear();
    	}
    }
    

    @Override
    public void rollback(ICompensation compensation) throws Exception {
        if (ShardTransactionStatusEnums.Rollback.equals(this.status)) {return;}
        this.status = ShardTransactionStatusEnums.Rollback;
        //  执行每一个任务的回滚
        List<TxResult> success = new ArrayList<>(),
                        failes = new ArrayList<>();
        boolean isCompensation = false;
        for (TransactionTask t : tasks) {
            try {
                //  回滚时，如果发现事务已经完整提交，走补偿
                if (t.ts.isCompleted()) {
                    log.warn(String.format("[ShardTransaction withTxAutoCommit] 执行事务回滚，但发现事务已经完整提交，走补偿策略 name:%s", t.name), t.e);
                    success.add(TxResult.build(t.name, true, t.res, null).metadata(t.metadata));
                    isCompensation = true;
                } else {
                    failes.add(TxResult.build(t.name, false, null, t.e).metadata(t.metadata));
                }
                transactionManager.rollback(t.ts);
            } catch (Exception e) {
                log.error(String.format("[ShardTransaction withTxAutoCommit] 事务回滚出现异常 name:%s", t.name), t.e);
            }
        }
        //  如果已经存在提交成功的子事务，执行补偿
        if (isCompensation && compensation != null) {
            compensation.doCompensation(success, failes);
        }
    }


    public static enum ShardTransactionStatusEnums {
        Running,
        Commited,
        Rollback;
    }
}
